#include "BasicBlock.hpp"
#include "Constant.hpp"
#include "Function.hpp"
#include "IRBuilder.hpp"
#include "Module.hpp"
#include "Type.hpp"

#include <iostream>
#include <memory>

// 定义宏，用于简化常量整数的创建
#define CONST_INT(num) ConstantInt::get(num, module.get())

// 生成 main 函数的 IR
void generateMainIR(std::unique_ptr<Module>& module) {
    auto builder = std::unique_ptr<IRBuilder>(new IRBuilder(nullptr, module.get()));
    Type* int32Type = module->get_int32_type();

    // 创建 main 函数
    auto mainFunction = Function::create(FunctionType::get(int32Type, {}), "main", module.get());
    if (!mainFunction) {
        std::cerr << "无法创建 main 函数" << std::endl;
        return;
    }

    // 创建入口块
    auto entryBasicBlock = BasicBlock::create(module.get(), "entry", mainFunction);
    if (!entryBasicBlock) {
        std::cerr << "无法创建入口基本块" << std::endl;
        return;
    }
    builder->set_insert_point(entryBasicBlock);

    // 分配变量 a 和 i，并初始化
    auto aAlloca = builder->create_alloca(int32Type);
    if (!aAlloca) {
        std::cerr << "无法为变量 a 分配内存" << std::endl;
        return;
    }
    auto iAlloca = builder->create_alloca(int32Type);
    if (!iAlloca) {
        std::cerr << "无法为变量 i 分配内存" << std::endl;
        return;
    }
    builder->create_store(CONST_INT(10), aAlloca); // 初始化 a = 10
    builder->create_store(CONST_INT(0), iAlloca);  // 初始化 i = 0

    // 创建条件判断块
    auto conditionBasicBlock = BasicBlock::create(module.get(), "cond", mainFunction);
    if (!conditionBasicBlock) {
        std::cerr << "无法创建条件判断基本块" << std::endl;
        return;
    }
    builder->create_br(conditionBasicBlock); // 跳转到条件判断块

    // 设置插入点到条件判断块
    builder->set_insert_point(conditionBasicBlock);
    auto iLoad = builder->create_load(iAlloca); // 加载 i 的值
    auto comparison = builder->create_icmp_lt(iLoad, CONST_INT(10)); // 比较 i < 10

    // 创建循环体块和结束块
    auto bodyBasicBlock = BasicBlock::create(module.get(), "body", mainFunction);
    if (!bodyBasicBlock) {
        std::cerr << "无法创建循环体基本块" << std::endl;
        return;
    }
    auto endBasicBlock = BasicBlock::create(module.get(), "end", mainFunction);
    if (!endBasicBlock) {
        std::cerr << "无法创建结束基本块" << std::endl;
        return;
    }
    builder->create_cond_br(comparison, bodyBasicBlock, endBasicBlock); // 根据条件跳转

    // 设置插入点到循环体块
    builder->set_insert_point(bodyBasicBlock);
    auto iLoadInBody = builder->create_load(iAlloca); // 加载 i 的值
    auto incrementedI = builder->create_iadd(iLoadInBody, CONST_INT(1)); // 计算 i + 1
    builder->create_store(incrementedI, iAlloca); // 更新 i 的值
    auto aLoad = builder->create_load(aAlloca); // 加载 a 的值
    auto updatedA = builder->create_iadd(aLoad, incrementedI); // 计算 a + (i + 1)
    builder->create_store(updatedA, aAlloca); // 更新 a 的值
    builder->create_br(conditionBasicBlock); // 跳回条件判断块

    // 设置插入点到结束块
    builder->set_insert_point(endBasicBlock);
    auto returnValue = builder->create_load(aAlloca); // 加载 a 的最终值
    builder->create_ret(returnValue); // 返回 a 的值
}

int main() {
    // 创建模块
    auto module = std::unique_ptr<Module>(new Module());

    // 生成 main 函数的 IR
    generateMainIR(module);

    // 输出生成的 IR 指令
    std::cout << module->print();

    return 0;
}
